import logging
from typing import Iterable, Iterator

import grpc
import numpy as np
from google.protobuf import struct_pb2

from parnassus.envs.arm.gym_env import GymEnv
from parnassus.protos.arm_stream import arm_stream_pb2 as pb2
from parnassus.protos.arm_stream import arm_stream_pb2_grpc as pb2_grpc
from parnassus.utils.tensor import numpy_to_tensor, tensor_to_numpy

logger = logging.getLogger(__name__)


class ArmEnvServicer(pb2_grpc.ArmEnvServicer):
    """ArmEnv gRPC服务实现"""

    def __init__(self):
        self.env = None
        logger.info("ArmEnvServicer initialized")

    def StreamEnv(
        self,
        request_iterator: Iterable[pb2.EnvRequest],
        context: grpc.ServicerContext,
    ) -> Iterator[pb2.EnvReply]:
        """处理双向流请求"""
        logger.info("StreamEnv method called")

        try:
            for request in request_iterator:
                if request.HasField("reset"):
                    yield self._handle_reset(request.reset)
                elif request.HasField("step"):
                    yield self._handle_step(request.step)
                elif request.HasField("close"):
                    yield self._handle_close()
                    break
                else:
                    logger.warning("Unknown request type received")
        except Exception as exc:  # pragma: no cover - gRPC will surface the error
            logger.exception("Error in StreamEnv")
            context.set_code(grpc.StatusCode.INTERNAL)
            context.set_details(f"Internal server error: {exc}")

    def _handle_reset(self, reset_request: pb2.Reset) -> pb2.EnvReply:
        """处理reset请求"""
        try:
            if self.env is None:
                self.env = GymEnv()
                logger.info("Environment created")

            # 使用seed参数（如果提供）
            seed = reset_request.seed[0] if reset_request.seed else None
            obs, info = self.env.reset(seed=seed)

            logger.info(f"Environment reset with observation shape: {obs.shape}")

            return pb2.EnvReply(reset=pb2.ResetReply(observation=numpy_to_tensor(obs)))
        except Exception as exc:
            logger.exception("Error in reset")
            raise

    def _handle_step(self, step_request: pb2.Step) -> pb2.EnvReply:
        """处理step请求"""
        try:
            if self.env is None:
                raise RuntimeError("Environment not initialized. Call reset first.")

            # 解析action
            action = tensor_to_numpy(step_request.action)

            # 执行步骤
            obs, reward, terminated, truncated, info = self.env.step(action)

            logger.debug(
                f"Step executed: obs_shape={obs.shape}, reward={reward}, "
                f"terminated={terminated}, truncated={truncated}"
            )

            struct_info = struct_pb2.Struct()
            try:
                struct_info.update(
                    {
                        key: (
                            value.tolist()
                            if hasattr(value, "tolist")
                            else value
                        )
                        for key, value in dict(info).items()
                    }
                )
            except Exception:
                logger.debug("Failed to serialise info dict; sending empty Struct")

            return pb2.EnvReply(
                step=pb2.StepReply(
                    observation=numpy_to_tensor(np.asarray(obs, dtype=np.float32)),
                    reward=numpy_to_tensor(
                        np.asarray([reward], dtype=np.float32)
                    ),
                    terminated=numpy_to_tensor(
                        np.asarray([terminated], dtype=np.bool_)
                    ),
                    truncated=numpy_to_tensor(
                        np.asarray([truncated], dtype=np.bool_)
                    ),
                    info=struct_info,
                )
            )
        except Exception as exc:
            logger.exception("Error in step")
            raise

    def _handle_close(self) -> pb2.EnvReply:
        """处理close请求"""
        try:
            if self.env is not None:
                self.env.close()
                self.env = None
                logger.info("Environment closed")

            return pb2.EnvReply(close=pb2.CloseReply())
        except Exception as exc:
            logger.exception("Error in close")
            raise
